using Crux.Core.LangChain.Abstractions.Databases;
using Crux.Core.LangChain.Abstractions.Databases.Metadatas;
using Crux.Core.LangChain.Abstractions.DocumentLoaders;
using Crux.Core.LangChain.Abstractions.Embeddings;
using Crux.Core.LangChain.Abstractions.Splitters.Text;

namespace Crux.Core.LangChain.Abstractions.Extensions;
/// <summary>
/// 
/// </summary>
public static class VectorDatabaseExtensions
{
    /// <summary>
    /// 
    /// </summary>
    /// <typeparam name="TLoader"></typeparam>
    /// <param name="vectorDatabase"></param>
    /// <param name="service"></param>
    /// <param name="dataSource"></param>
    /// <param name="textSplitter"></param>
    /// <param name="collectionName"></param>
    /// <param name="settings"></param>
    /// <param name="dimensions"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    public static async Task<IVectorCollection> AddDocumentsFromAsync<TLoader>(
        this IVectorDatabase vectorDatabase,
        ITextEmbeddingGenerationService service,
        DataSource dataSource,
        ITextSplitter? textSplitter = null,
        string collectionName = "langchain",
        DocumentLoaderSettings? settings = null,
        int dimensions = 1536,
        CancellationToken cancellationToken = default
    ) where TLoader : IDocumentLoader
    {
        if (service is null)
        {
            throw new ArgumentNullException(nameof(service));
        }

        if (dataSource is null)
        {
            throw new ArgumentNullException(nameof(dataSource));
        }

        var splitDocuments = await AddSplitDocumentsAsync<TLoader>(dataSource, textSplitter, settings, cancellationToken).ConfigureAwait(false);

        var vectorCollection = await vectorDatabase.GetCollectionAsync(collectionName, dimensions, cancellationToken).ConfigureAwait(false);

        var vectors = await GeneratorVectorsAsync(service, dimensions, splitDocuments, cancellationToken).ConfigureAwait(false);

        await vectorCollection.AddAsync(vectors, cancellationToken).ConfigureAwait(false);

        return vectorCollection;
    }

    private static async Task<IReadOnlyList<Vector>> GeneratorVectorsAsync(
        ITextEmbeddingGenerationService embeddingGenerationService,
        int dimensions,
        IReadOnlyList<Document> splitDocuments,
        CancellationToken cancellationToken)
    {
        var texts = splitDocuments.Select(doc => doc.PageContent).ToArray();
        var metadatas = splitDocuments.Select(doc => doc.Metadata).ToArray();

        var embeddings = await embeddingGenerationService.GenerateEmbeddingsAsync(texts, dimensions, cancellationToken).ConfigureAwait(false);

        var vectors = texts.Select((text, i) => new Vector
        {
            Text = text,
            Metadata = metadatas?.ElementAt(i).ToDictionary(x => x.Key, x => x.Value),
            Embedding = embeddings[i],
        }).ToArray();

        return vectors;
    }

    private static async Task<IReadOnlyList<Document>> AddSplitDocumentsAsync<TLoader>(
        DataSource dataSource,
        ITextSplitter? textSplitter,
        DocumentLoaderSettings? settings,
        CancellationToken cancellationToken) where TLoader : IDocumentLoader
    {
        var loader = Activator.CreateInstance<TLoader>();

        var documents = await loader.LoadAsync(dataSource, settings, cancellationToken).ConfigureAwait(false);

        textSplitter ??= new CharacterTextSplitter();

        return textSplitter.SplitDocuments(documents);
    }
}